"""添加ConversationSession表的数据库迁移脚本"""

import sys
import os

# 添加项目根目录到Python路径
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

from flask import Flask
from extensions import db
from models import ConversationSession


def create_conversation_sessions_table():
    """创建conversation_sessions表"""
    try:
        # 创建表
        db.create_all()
        print("✓ ConversationSession表创建成功")

        # 检查表是否存在
        inspector = db.inspect(db.engine)
        tables = inspector.get_table_names()

        if 'conversation_sessions' in tables:
            print("✓ conversation_sessions表已存在")

            # 显示表结构
            columns = inspector.get_columns('conversation_sessions')
            print("\n表结构:")
            for col in columns:
                print(f"  - {col['name']}: {col['type']}")
        else:
            print("✗ conversation_sessions表创建失败")

    except Exception as e:
        print(f"✗ 创建表时出错: {e}")
        raise


def migrate_existing_data():
    """迁移现有的ConversationFile数据到ConversationSession"""
    try:
        # 检查是否存在旧表
        inspector = db.inspect(db.engine)
        tables = inspector.get_table_names()

        if 'conversation_files' not in tables:
            print("✓ 没有需要迁移的旧数据")
            return

        # 执行数据迁移
        migration_sql = """
                        INSERT INTO conversation_sessions (user_id, session_id, title, conversation_dir,
                                                           timestamp_folder,
                                                           message_count, total_tokens, total_cost, conversation_file,
                                                           metadata_file,
                                                           total_size, created_at, updated_at, last_message_at)
                        SELECT user_id,
                               CAST(request_id AS TEXT)                    as session_id,
                               COALESCE(query_summary, 'Migrated Session') as title,
                               conversation_dir,
                               timestamp_folder,
                               1                                           as message_count,
                               0                                           as total_tokens,
                               0.0                                         as total_cost,
                               conversation_dir || '/conversation.jsonl'   as conversation_file,
                               conversation_dir || '/metadata.json'        as metadata_file,
                               COALESCE(total_size, 0)                     as total_size,
                               created_at,
                               updated_at,
                               updated_at                                  as last_message_at
                        FROM conversation_files
                        WHERE request_id NOT IN (SELECT CAST(session_id AS INTEGER)
                                                 FROM conversation_sessions
                                                 WHERE session_id ~ '^[0-9]+$'
                            ); \
                        """

        db.session.execute(db.text(migration_sql))
        db.session.commit()

        print("✓ 数据迁移完成")

    except Exception as e:
        print(f"✗ 数据迁移失败: {e}")
        db.session.rollback()
        raise


def run_migration():
    """运行完整的迁移过程"""
    print("开始数据库迁移...")

    # 1. 创建新表
    create_conversation_sessions_table()

    # 2. 迁移现有数据
    migrate_existing_data()

    print("\n✓ 数据库迁移完成！")
    print("\n新的对话记录系统特性:")
    print("  - 基于session_id组织对话")
    print("  - 同一会话的多轮对话记录在一起")
    print("  - 使用JSONL格式存储对话历史")
    print("  - 自动生成会话标题")
    print("  - 统计会话级别的token和费用")


if __name__ == '__main__':
    # 创建Flask应用上下文
    app = Flask(__name__)
    app.config.from_object('config.Config')

    # 初始化数据库
    db.init_app(app)

    with app.app_context():
        run_migration()
